/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package idv.cloudee.proxy2.impl.goodies;

import com.cirnoworks.http.utils.Header;
import com.cirnoworks.http.utils.HeaderRequest;
import com.cirnoworks.http.utils.HeaderResponse;
import idv.cloudee.proxy2.framework.ProxySession;
import idv.cloudee.proxy2.impl.getpost.GETPOSTModifier;
import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Map.Entry;

/**
 *
 * @author yuxuanhuang
 */
public class GETPOSTStat implements GETPOSTModifier {

    private GETPOSTStatFactory factory;

    public GETPOSTStat(GETPOSTStatFactory factory) {
        this.factory = factory;
    }

    public void beginRequest(ProxySession session) {
    }

    public HeaderRequest modifyRequest(HeaderRequest param) {
        factory.addUpload(param.getRequestLine().length() + 2);
        factory.addUpload(countHeader(param));
        return param;
    }

    public OutputStream getPOSTOutputStream(OutputStream os) {
        return new CountUploadOutputStream(os);
    }

    public HeaderResponse modifyResponse(HeaderResponse param) {
        factory.addDownload(param.getResponseLine().length() + 2);
        factory.addDownload(countHeader(param));
        return param;
    }

    public OutputStream getResponseOutputStream(OutputStream os) {
        return new CountDownloadOutputStream(os);
    }

    public void requestOver() {
    }

    private int countHeader(Header param) {
        int len = 0;
        for (Entry<String, ArrayList<String>> entries : param.entrySet()) {
            String key = entries.getKey();
            for (String value : entries.getValue()) {
                len += key.length() + value.length() + 4;
            }
        }
        return len + 2;
    }

    class CountUploadOutputStream extends OutputStream {

        private OutputStream os;

        public CountUploadOutputStream(OutputStream os) {
            this.os = os;
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            factory.addUpload(len);
            os.write(b, off, len);
        }

        @Override
        public void write(byte[] b) throws IOException {
            factory.addUpload(b.length);
            os.write(b);
        }

        public void write(int b) throws IOException {
            factory.addUpload(1);
            os.write(b);
        }

        @Override
        public void flush() throws IOException {
            os.flush();
        }

        @Override
        public void close() throws IOException {
            os.close();
        }
    }

    class CountDownloadOutputStream extends OutputStream {

        private OutputStream os;

        public CountDownloadOutputStream(OutputStream os) {
            this.os = os;
        }

        @Override
        public void write(byte[] b, int off, int len) throws IOException {
            factory.addDownload(len);
            os.write(b, off, len);
        }

        @Override
        public void write(byte[] b) throws IOException {
            factory.addDownload(b.length);
            os.write(b);
        }

        public void write(int b) throws IOException {
            factory.addDownload(1);
            os.write(b);
        }

        @Override
        public void flush() throws IOException {
            os.flush();
        }

        @Override
        public void close() throws IOException {
            os.close();
        }
    }
}
